package generator

import (
	"bytes"
	"github.com/golang/protobuf/proto"
	"google.golang.org/protobuf/types/pluginpb"
	"io"
	"os"
)

func New() *Generator {
	g := new(Generator)
	g.Buf = new(bytes.Buffer)
	g.Req = new(pluginpb.CodeGeneratorRequest)
	g.Rsp = new(pluginpb.CodeGeneratorResponse)
	return g
}

type Generator struct {
	Buf *bytes.Buffer
	Req *pluginpb.CodeGeneratorRequest  // The input.
	Rsp *pluginpb.CodeGeneratorResponse // The output.
}

func (g *Generator) Run(f func(*Generator)) error {
	err := g.read()
	if err != nil {
		return err
	}
	f(g)
	out, err := proto.Marshal(g.Rsp)
	if err != nil {
		return err
	}
	if _, err := os.Stdout.Write(out); err != nil {
		return err
	}
	return nil
}

func (g *Generator) read() error {
	data, err := io.ReadAll(os.Stdin)
	if err != nil {
		return err
	}
	g.Req.Reset()
	g.Buf.Write(data)
	return proto.Unmarshal(data, g.Req)
}

func (g *Generator) RunTest(data []byte, f func(*Generator)) error {
	err := g.readTest(data)
	if err != nil {
		return err
	}
	f(g)
	out, err := proto.Marshal(g.Rsp)
	if err != nil {
		return err
	}
	if _, err := os.Stdout.Write(out); err != nil {
		return err
	}
	return nil
}

func (g *Generator) readTest(data []byte) error {
	g.Req.Reset()
	g.Buf.Write(data)
	return proto.Unmarshal(data, g.Req)
}
